//
// Created by haoy on 2017/6/20.
//

#include "ulpfec_decoder.hpp"

#include <algorithm>
#include <list>

#include "codecs.hpp"
#include "xcutil.h"
#include "xm_log.h"
#include "xrtc_endpoint.h"

// todo: 通过ULPFEC还原出来的RTP包EXT部分有错误，
// 查看FEC的PAYLOAD发现EXT对应的几个字节为：BE DE 00 01 32 00 00 00
// 怀疑是生成FEC时，ABS-SEND-TIME的值是00 00 00
// 待确认

UlpfecDecoder::UlpfecDecoder(RtpFactory* factory, rtc_codecs_t* codecs)
        : rtp_factory_(factory), codecs_(codecs) {}

UlpfecDecoder::~UlpfecDecoder() {
    media_packets_.clear();
    fec_packets_.clear();
}

void UlpfecDecoder::add(std::shared_ptr<RtpPacket> packet) {
    uint16_t seq = packet->header.seq;
    if (tail_seq_ == head_seq_ && tail_seq_ == 0) {
        tail_seq_ = seq;
        head_seq_ = seq - WINDOW_SIZE;
    } else if ((head_seq_ < tail_seq_ && (seq < head_seq_ || seq > tail_seq_))
        || (head_seq_ > seq && seq > tail_seq_)) {
        uint16_t udelta = seq - tail_seq_;
        if (udelta < 3000) {
            tail_seq_ = seq;
            head_seq_ = seq - WINDOW_SIZE;
        }
    }

    if (packet->codec_id == rtc_codec_id_ulpfec) {
        if (fec_packets_.count(packet->header.seq) != 0) {
            return;
        }
        fec_packets_.emplace(packet->header.seq, packet);
    } else {
        if (media_packets_.count(packet->header.seq) != 0) {
            return;
        }
        media_packets_.emplace(packet->header.seq, packet);
    }

    char media_packet_str[2048] = {0,};
    int media_packet_str_len = 0;

    // remove packet out of window
    auto mit = media_packets_.begin();
    while (mit != media_packets_.end()) {
        uint16_t mseq = mit->first;
        if (head_seq_ < tail_seq_) {
            if (mseq < head_seq_ || mseq > tail_seq_) {
                if (debug_) {
                    media_packet_str_len += sprintf(media_packet_str + media_packet_str_len, "X%s%u,", mit->second->is_recovered ? "r" : "", mseq);
                }
                media_packets_.erase(mit++);
            } else {
                if (debug_) {
                    media_packet_str_len += sprintf(media_packet_str + media_packet_str_len, "%s%u,", mit->second->is_recovered ? "r" : "", mseq);
                }
                mit++;
            }
        } else {
            if (mseq > tail_seq_ && mseq < head_seq_) {
                media_packet_str_len += sprintf(media_packet_str + media_packet_str_len, "X%s%u,", mit->second->is_recovered ? "r" : "", mseq);
                media_packets_.erase(mit++);
            } else {
                media_packet_str_len += sprintf(media_packet_str + media_packet_str_len, "%s%u,", mit->second->is_recovered ? "r" : "", mseq);
                mit++;
            }
        }
    }

    if (debug_) {
        dbgi("fec media packets: [%s]", media_packet_str);
        memset(media_packet_str, 0, media_packet_str_len);
        media_packet_str_len = 0;
    }

    auto fit = fec_packets_.begin();
    while (fit != fec_packets_.end()) {
        auto& fec = fit->second->fec;
        uint16_t fseq = fec.header.seq;
        if (head_seq_ < tail_seq_) {
            if (fseq < head_seq_ || fseq > tail_seq_) {
                if (debug_) {
                    media_packet_str_len += sprintf(media_packet_str + media_packet_str_len, "X%u,", fit->second->header.seq);
                }
                fec_packets_.erase(fit++);
            } else {
                if (debug_) {
                    media_packet_str_len += sprintf(media_packet_str + media_packet_str_len, "%u,", fit->second->header.seq);
                }
                ++fit;
            }
        } else {
            if (fseq > tail_seq_ && fseq < head_seq_) {
                if (debug_) {
                    media_packet_str_len += sprintf(media_packet_str + media_packet_str_len, "X%u,", fit->second->header.seq);
                }
                fec_packets_.erase(fit++);
            } else {
                if (debug_) {
                    media_packet_str_len += sprintf(media_packet_str + media_packet_str_len, "X%u,", fit->second->header.seq);
                }
                ++fit;
            }
        }
    }

    if (debug_) {
        dbgi("fec fec packets: [%s]", media_packet_str);
    }

}

void dump_rtp_extensions(const uint8_t *buf, size_t len) {
    const uint8_t* ext_read = buf + 4;
    const uint8_t* ext_end = buf + len;
    //rtc_dbgi("\t - exten length: %u", len);
    if (buf[0] == 0xBE && buf[1] == 0xDE) {
        while (ext_read < ext_end) {
            uint8_t id = ext_read[0] >> 4;
            if (id == 0) {
                ext_read++;
                continue;
            }
            if (id == 15) {
                break;
            }
            uint8_t l = ext_read[0] & 0x0F;
            static char str[32+1];
            uint8_t i = 0;
            while (i <= l+1) {
                sprintf(str + 2 * i, "%02X", ext_read[1+i]);
                i++;
            }
            str[(l+1)*2] = 0;
            dbgi("\t - exten(1 byte) id=%d l=%d value=0x%s", id, l, str);
            ext_read += l + 1 + 1;
        }
    } else if (buf[0] == 0x10 && buf[1] == 0x00) {
        while (ext_read < ext_end) {
            uint8_t id = ext_read[0];
            if (id == 0) {
                ext_read++;
                continue;
            }
            if (id == 15) {
                break;
            }
            uint8_t l = ext_read[1];
            if (l > 0) {
                static char str[512+1];
                uint8_t i = 0;
                while (i < l) {
                    sprintf(str + 2 * i, "%X", ext_read[2+i]);
                    i++;
                }
                str[l*2+1] = 0;
                dbgi("\t - exten(2 byte) id=%d l=%d value=0x%s", id, l, str);
            } else {
                dbgi("\t - exten(2 byte) id=%d l=%d", id, l);
            }
            ext_read += l + 2;
        }
    }
}

static void hex_print(const char * prefix, const unsigned char * data, int length){
    static char str_buf[65535] = {0,};
    char* p = str_buf;
    int i = 0;

    for(i = 0; i < length; ++i) {
        p += sprintf(p, "%.2X:", data[i]);
    }
    *(p-1) = 0;

    dbgi("%s %s", prefix, str_buf);
}

static
void dump_rtp(const char * prefix, void* buf_, int len ){
    const unsigned char * buf = (const unsigned char *)buf_;

    unsigned char v =  (buf[0]>>6) & 0x3;
    unsigned char p =  (buf[0]>>5) & 0x1;
    unsigned char x =  (buf[0]>>4) & 0x1;
    unsigned char cc = (buf[0]>>0) & 0xF;
    unsigned char m =  (buf[1]>>7) & 0x1;
    unsigned char pt = (buf[1]>>0) & 0x7F;
    unsigned short seq = be_get_u16(buf+2);
    unsigned int ts = be_get_u32(buf+4);
    unsigned int ssrc = be_get_u32(buf+8);

    int header_len = 12 + cc * 4;
    int ext_len = 0;
    if (x && len > header_len + 4) {
        ext_len = be_get_u16(buf+header_len+2);
        header_len += 4 + 4 * ext_len;
    }

    int payload_len = len - header_len;
    dbgi("fec %s rtp[v=%u,p=%u,x=%u,cc=%u,m=%u,pt=%u,seq=%u,ts=%u,ssrc=%u,pl=%d,len=%d]", prefix,
            v, p, x, cc, m, pt, seq, ts, ssrc, payload_len, len);

    if (x == 1) {
        dump_rtp_extensions(buf + 12 + 4 * cc, 4 + 4 * ext_len);
    }

    hex_print("total", buf, std::min(64, len));
}

// todo: 当前用64位bitset匹配media/fec，最大恢复64个包之内的包，支持到3000个
int UlpfecDecoder::recover(std::list<std::shared_ptr<RtpPacket>>& recovered_packets) {
    recovered_packets.clear();
    if (media_packets_.empty() || fec_packets_.empty()) {
        return 0;
    }
    // todo: check if we can simply start from media_packets_.begin()
    auto mit = media_packets_.lower_bound(head_seq_);
    // auto first_media_seq = mit->first;
    uint64_t media_mask = 0;
    uint8_t media_mask_bits_available = 64;
    std::list<uint16_t> lost_seqs;
    // initial zeroes
    for (auto si = head_seq_; si < mit->first && media_mask_bits_available; ++si, --media_mask_bits_available) {
        media_mask = media_mask << 1;
        media_mask &= 0xFFFFFFFFFFFFFFFE;
    }
    uint16_t next_seq = mit->first + 1;
    while (media_mask_bits_available && mit != media_packets_.end()) {
        uint16_t seq = mit->first;
        // insert zeroes for each packet lost
        for (auto si = next_seq; si < seq && media_mask_bits_available; ++si, --media_mask_bits_available) {
            media_mask = media_mask << 1;
            media_mask &= 0xFFFFFFFFFFFFFFFE;
            if (fec_packets_.count(si) == 0) {
                lost_seqs.emplace_back(si);
            }
        }
        if (media_mask_bits_available == 0) {
            break;
        }
        // insert one for received media packet
        media_mask = media_mask << 1;
        media_mask |= 1;
        --media_mask_bits_available;
        next_seq = seq + 1;
        ++mit;
    }
    // fill remains zeroes
    while (media_mask_bits_available) {
        media_mask = media_mask << 1;
        media_mask &= 0xFFFFFFFFFFFFFFFE;
        --media_mask_bits_available;
    }

    if (lost_seqs.empty()) {
        return 0;
    }

    // for each lost packet
    for (auto lost_seq : lost_seqs) {
        for (auto& fit : fec_packets_) {
            auto fec = fit.second;
            std::shared_ptr<RtpPacket> rtp = fit.second;
            bool lost_packet_protected =
                    std::find(fec->fec.level_header.protected_seqs.begin(),
                            fec->fec.level_header.protected_seqs.end(),
                            lost_seq) != fec->fec.level_header.protected_seqs.end();
            if (!lost_packet_protected) {
                continue;
            }
            // 1 means first bit
            uint16_t lost_bit_in_fec_mask = lost_seq - fec->fec.header.seq + 1;
            if (fec->fec.header.seq < head_seq_) {
                dbgw("fec: bug caught, fec packet %u has sn base smaller than head, %u < %u",
                        rtp->header.seq, fec->fec.header.seq, head_seq_);
                continue;
            }
            uint64_t aligned_media_mask = media_mask << (fec->fec.header.seq - head_seq_);
            uint64_t fec_mask_except_lost = fec->fec.level_header.mask & ~(((uint64_t)1) << (64 - lost_bit_in_fec_mask));
            bool got_all_needed_media_packets = (fec_mask_except_lost & aligned_media_mask) == fec_mask_except_lost;
            if (got_all_needed_media_packets) {
                char media_packets_seqs[1024] = {0,};
                int media_packets_seqs_len = 0;
                for (auto id : fec->fec.level_header.protected_seqs) {
                    media_packets_seqs_len += sprintf(media_packets_seqs + media_packets_seqs_len, "%d,", id);
                }
                if (debug_) {
                    dbgi("fec: going to recover %d from media [%s] and fec %d",
                            lost_seq, media_packets_seqs, fec->header.seq);
                }
                static uint8_t recover_buffer[RtpPacket::MAX_BUF];
                memset(recover_buffer, 0, RtpPacket::MAX_BUF);
                int v = 2;
                int p = fec->fec.header.p;
                int x = fec->fec.header.x;
                int cc = fec->fec.header.cc;
                recover_buffer[0] = v << 6 | p << 5 | x << 4 | cc;
                recover_buffer[1] = fec->fec.header.m << 7 | fec->fec.header.pt;
                be_set_u16(lost_seq, recover_buffer + 2);
                be_set_u32(fec->fec.header.ts, recover_buffer + 4);
                be_set_u32(rtp->header.ssrc, recover_buffer + 8);
                int protection_size_0 = fec->fec.level_header.protection_len;
                uint16_t length_recovery = fec->fec.header.len;

                int n = 12;
                for (int i = 0; i < protection_size_0; i++) {
                    recover_buffer[n++] = fec->fec.payload[i];
                }
                if (debug_) {
                    hex_print("fec: fec payload", fec->fec.payload,
                            std::min(fec->fec.payload_size, 32));
                }
                uint16_t media_seq = fec->fec.header.seq;
                while (fec_mask_except_lost > 0) {
                    uint8_t msb = fec_mask_except_lost >> 63;
                    if (msb > 0) {
                        auto media_packet = media_packets_[media_seq];
                        // 5.  Set the version of the new packet to 2.  Skip the first 2 bits
                        // in the recovery bit string.
                        recover_buffer[0] ^= media_packet->buffer[0] & 0x3F;
                        recover_buffer[1] ^= media_packet->buffer[1];
                        // byte 2-3 are seq number, do not xor it
                        recover_buffer[4] ^= media_packet->buffer[4];
                        recover_buffer[5] ^= media_packet->buffer[5];
                        recover_buffer[6] ^= media_packet->buffer[6];
                        recover_buffer[7] ^= media_packet->buffer[7];
                        // byte 8-11 are ssrc, do not xor it
                        int limit = std::min(media_packet->buffer_size, 12+protection_size_0);
                        for (auto n = 12; n < limit; n++) {
                            recover_buffer[n] ^= media_packet->buffer[n];
                        }

                        length_recovery ^= (media_packet->payload_size
                                + media_packet->csrc_len + media_packet->ext_len);
                        if (debug_) {
                            if (media_packet->is_recovered) {
                                dbgi("fec: xor packet %d(recovered)", media_packet->header.seq);
                            } else {
                                dbgi("fec: xor packet %d", media_packet->header.seq);
                            }
                            dump_rtp("fec: media packet", media_packet->buffer, media_packet->buffer_size);
                        }
                    }
                    fec_mask_except_lost = fec_mask_except_lost << 1;
                    media_seq++;
                }

                if (length_recovery > RtpPacket::MAX_BUF || length_recovery < 20 || length_recovery > protection_size_0) {
                    dbge("fec: recovered ridiculous payload length: %d", length_recovery);
                    continue;
                }

                auto recovered_packet = rtp_factory_->allocate();
                int recovered_packet_size = recovered_packet->read(recover_buffer, 12 + length_recovery, codecs_);
                if (recovered_packet_size == 0) {
                    continue;
                } else {
                    recovered_packet->is_recovered = true;
                    recovered_packets.emplace_back(recovered_packet);
                }

                auto saved_recovered_packet = rtp_factory_->allocate();
                if (saved_recovered_packet->read(recover_buffer,
                        12 + length_recovery, codecs_) == 0) {
                    continue;
                } else {
                    saved_recovered_packet->is_recovered = true;
                    media_packets_.emplace(lost_seq, saved_recovered_packet);
                    uint16_t lost_bit_in_media_mask = lost_seq - head_seq_ + 1;
                    media_mask ^= ((uint64_t)1) << (64 - lost_bit_in_media_mask);
                }
            } // end of if (got_all_needed_media_packets)
        } // end of for each fec packet
    } // end of for each lost media packet

    return recovered_packets.size();
}